// 引入标准输入输出库（用于printf等函数）
#include <cstdio>        
// 引入标准库函数（malloc, free, drand48等）
#include <cstdlib>       
// 引入CUDA运行时库头文件（用于GPU相关操作）
#include <cuda_runtime.h> 

// 定义矩阵访问宏（行优先存储）
// 功能：将二维索引(i,j)转换为行优先存储的一维数组下标
// 参数：i-行号，j-列号，a-矩阵数据指针，n-矩阵列数
#define A(i, j) a[(i)*n + (j)] // 计算下标公式：行号×总列数 + 列号

// CUDA核函数声明
// 功能：矩阵乘法 C = A * B（单精度浮点）
// 参数：
//   A-输入矩阵1（设备内存指针）
//   B-输入矩阵2（设备内存指针）
//   C-输出矩阵（设备内存指针）
//   M-矩阵C的行数（等于矩阵A的行数）
//   N-矩阵C的列数（等于矩阵B的列数）
//   K-矩阵A的列数/矩阵B的行数（中间维度）
// __global__ void cuda_sgemm<BLOCK,k>(float* A, float* B, float *C, 
//                           int M, int N, int K);

// 生成随机矩阵函数
// 参数：
//   m-矩阵行数
//   n-矩阵列数
//   a-矩阵数据指针（主机内存）
void random_matrix(int m, int n, float* a) {
    // 双重循环遍历矩阵所有元素
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
// 条件编译：选择初始化方式
#if 1  // 使用随机数初始化（-1~1之间的均匀分布）
            // drand48()生成[0,1)的随机数，通过公式映射到[-1,1)
            //A(i, j) = 2.0f * (float)drand48() - 1.0f; 
            A(i, j) =i*n+j; 
#else  // 使用固定模式初始化（调试用）
            // 生成固定模式数据：(列号 - 行号)取模3
            A(i, j) = (j - i) % 3;
#endif
        }
    }
}

// CPU矩阵乘法（朴素实现）
// 参数：
//   A_ptr-输入矩阵1（主机内存指针）
//   B_ptr-输入矩阵2（主机内存指针）
//   C_ptr-输出矩阵（主机内存指针）
//   M-输出矩阵行数
//   N-输出矩阵列数
//   K-中间维度（A的列数/B的行数）
void cpu_sgemm(float* A_ptr, float* B_ptr, float* C_ptr,
              const int M, const int N, const int K) 
{
    // 三重循环实现矩阵乘法
    for (int m = 0; m < M; m++) {        // 遍历输出矩阵的每一行
        for (int n = 0; n < N; n++) {    // 遍历输出矩阵的每一列
            float temp = 0.0f;           // 临时累加器
            for (int k = 0; k < K; k++) { // 内积计算循环
                // 计算A的第m行第k列元素 × B的第k行第n列元素
                // A_ptr[m*K + k] 对应A[m][k]
                // B_ptr[k*N + n] 对应B[k][n]
                temp += A_ptr[m * K + k] * B_ptr[k * N + n];
            }
            C_ptr[m * N + n] = temp;    // 存储计算结果到C[m][n]
        }
    }
}

// 矩阵比较函数
// 参数：
//   m-矩阵行数
//   n-矩阵列数
//   a-矩阵1数据指针
//   b-矩阵2数据指针
// 返回值：最大绝对差值
float compare_matrices(int m, int n, float *a, float *b) {
    int i, j;                          // 循环变量
    float max_diff = 0.0f, diff;       // 最大差值和临时差值
    int printed = 0;                   // 标记是否已打印第一个显著差异
    
    // 遍历所有矩阵元素
    for (i = 0; i < m; i++) {
        for (j = 0; j < n; j++) {
            // 计算绝对差值
            diff = fabsf(a[i * n + j] - b[i * n + j]);
            
            // 更新最大差值
            if (diff > max_diff) max_diff = diff;
            
            // 打印第一个超过阈值（0.5）的差异
            if (!printed && diff > 0.5f) {
                // 输出格式：位置、差值、两个矩阵的对应值
                printf("\nError: i=%d j=%d | diff=%.6f | A=%.6f B=%.6f", 
                      i, j, diff, a[i * n + j], b[i * n + j]);
                printed = 1; // 仅打印第一个显著差异
            }
        }
    }
    return max_diff; // 返回最大绝对差值
}

/*---------- CUDA核函数实现（全局内存版本） ----------*/
// 参数说明同声明部分
template <unsigned int BLOCK_SIZE>
__global__ void cuda_sgemm(float *A_ptr, float *B_ptr, float *C_ptr,int M, int N, int K) {
    // 计算当前线程处理的矩阵位置（全局索引）
    // blockIdx.y：当前块在网格中的Y坐标
    // blockDim.y：块在Y方向的维度（16）
    // threadIdx.y：线程在块内的Y坐标
    int row = blockIdx.y * blockDim.y + threadIdx.y; // 计算全局行号
    
    // blockIdx.x：当前块在网格中的X坐标
    // blockDim.x：块在X方向的维度（16）
    // threadIdx.x：线程在块内的X坐标
    int col = blockIdx.x * blockDim.x + threadIdx.x; // 计算全局列号

    // 计算当前线程块对应的矩阵A和B的起始位置
    float *A_ptr_start = A_ptr + blockIdx.y * blockDim.y * K; // A的起始行（每个块处理BLOCK行）

    float *B_ptr_start = B_ptr + blockIdx.x * blockDim.x;     // B的起始列（每个块处理BLOCK列）

    // 声明共享内存，用于缓存A和B的数据块
    __shared__ float a_shared[BLOCK_SIZE][BLOCK_SIZE];  // BLOCK_SIZE行，K_列（K_=K）
    __shared__ float b_shared[BLOCK_SIZE][BLOCK_SIZE];  // K_行，BLOCK_SIZE列（K_=K）
    float sum = 0.0f; // 累加器初始化 
    // 将A和B的全局内存数据分块加载到共享内存
    for(int s=0; s<K; s+=blockDim.x) {  // 按块大小步进，遍历K维度
        // 将A的当前块加载到a_shared
        a_shared[threadIdx.y][threadIdx.x] = A_ptr_start[threadIdx.x + s + threadIdx.y*K];
        printf( "a%i数据拷贝：%f \n", s,a_shared[threadIdx.y][threadIdx.x]);
        // 将B的当前块加载到b_shared
        b_shared[threadIdx.y][threadIdx.x] = B_ptr_start[threadIdx.x + (threadIdx.y + s)*N];
        printf( "b%i数据拷贝：%f \n", s,b_shared[threadIdx.y][threadIdx.x]);

        __syncthreads();
        

        for (int k = 0; k < BLOCK_SIZE; k++) {
            sum += a_shared[threadIdx.y][k] * b_shared[k][threadIdx.x];
        }

        __syncthreads();



    }
    
    

    
        
        
       
        // 将结果写入C矩阵的对应位置
     C_ptr[row * N + col] = sum;
    
}

// 主函数
int main() {
    // 矩阵维度配置（实际为512x512）
    int m = 8; // 矩阵A的行数，矩阵C的行数
    int n = 4; // 矩阵B的列数，矩阵C的列数
    int k = 4; // 矩阵A的列数，矩阵B的行数

    // 计算各矩阵需要的内存大小（字节数）
    const size_t mem_size_A = m * k * sizeof(float); // A矩阵：m行k列
    const size_t mem_size_B = k * n * sizeof(float); // B矩阵：k行n列
    const size_t mem_size_C = m * n * sizeof(float); // C矩阵：m行n列

    /*---------- 主机（CPU）内存分配 ----------*/
    float *matrix_A_host = (float*)malloc(mem_size_A);   // 分配A矩阵主机内存
    float *matrix_B_host = (float*)malloc(mem_size_B);   // 分配B矩阵主机内存
    float *matrix_C_host_gpu = (float*)malloc(mem_size_C); // 分配GPU结果存储空间
    float *matrix_C_host_cpu = (float*)malloc(mem_size_C); // 分配CPU结果存储空间

    /*---------- 设备（GPU）内存分配 ----------*/
    float *matrix_A_device, *matrix_B_device, *matrix_C_device;
    cudaMalloc((void **)&matrix_A_device, mem_size_A); // 分配设备内存给A
    cudaMalloc((void **)&matrix_B_device, mem_size_B); // 分配设备内存给B
    cudaMalloc((void **)&matrix_C_device, mem_size_C); // 分配设备内存给C

    // 初始化矩阵数据
    random_matrix(m, k, matrix_A_host); // 生成随机矩阵A（m行k列）
    random_matrix(k, n, matrix_B_host); // 生成随机矩阵B（k行n列）

    /*---------- 数据传输到GPU ----------*/
    // 将矩阵A从主机内存拷贝到设备内存
    cudaMemcpy(matrix_A_device, matrix_A_host, mem_size_A, cudaMemcpyHostToDevice);
    // 将矩阵B从主机内存拷贝到设备内存
    cudaMemcpy(matrix_B_device, matrix_B_host, mem_size_B, cudaMemcpyHostToDevice);

    /*---------- CPU计算 ----------*/
    cpu_sgemm(matrix_A_host, matrix_B_host, matrix_C_host_cpu, m, n, k);

    /*---------- GPU计算 ----------*/
    constexpr int BLOCK = 2;              // 线程块维度（2x2线程）
    dim3 block(BLOCK, BLOCK);              // 定义线程块结构（二维）
    dim3 grid((n + BLOCK - 1)/BLOCK,       // 计算网格X维度（向上取整）
              (m + BLOCK - 1)/BLOCK);      // 计算网格Y维度（向上取整）
    
    // 调用CUDA核函数，使用模板参数BLOCK和k
    cuda_sgemm<BLOCK><<<grid, block>>>(matrix_A_device, matrix_B_device,  matrix_C_device, m, n, k);

    /*---------- 将GPU结果拷贝回主机 ----------*/
    cudaMemcpy(matrix_C_host_gpu, matrix_C_device, mem_size_C, cudaMemcpyDeviceToHost);

    /*---------- 结果验证 ----------*/
    float diff = compare_matrices(m, n, matrix_C_host_gpu, matrix_C_host_cpu);
    printf("\n V2 sgemm diff: %f\n", diff); // 输出最大差值

    /*---------- 释放所有内存 ----------*/
    // 释放主机内存
    free(matrix_A_host);
    free(matrix_B_host);
    free(matrix_C_host_gpu);
    free(matrix_C_host_cpu);
    // 释放设备内存
    cudaFree(matrix_A_device);
    cudaFree(matrix_B_device);
    cudaFree(matrix_C_device);

    return 0;
}